package org.zoomdev.zoom.aop.modules;

import org.objectweb.asm.ClassVisitor;
import org.objectweb.asm.MethodVisitor;
import org.objectweb.asm.Opcodes;
import org.zoomdev.zoom.aop.AopFactory;
import org.zoomdev.zoom.aop.utils.AsmUtils;
import org.zoomdev.zoom.common.exceptions.ZoomException;

import java.lang.reflect.Method;
import java.util.ArrayDeque;
import java.util.ArrayList;
import java.util.Deque;
import java.util.List;
import java.util.Map;

public final class ReturnNewRewriter extends ClassVisitor {

    public final List<Method> modified = new ArrayList<>();
    private final AopFactory aop; // 原始类名 -> 子类名（内部格式，如 A -> SubA）
    private final Map<String, Method> map;

    public ReturnNewRewriter(ClassVisitor cv, AopFactory aop, Map<String, Method> map) {
        super(Opcodes.ASM9, cv);
        this.aop = aop;
        this.map = map;
    }

    /**
     * NEW A
     * DUP
     * NEW B
     * DUP
     * INVOKESPECIAL B.<init>()V
     * NEW C
     * DUP
     * INVOKESPECIAL C.<init>()V
     * INVOKESPECIAL A.<init>(LB;LC;)V
     * ARETURN
     *
     * @param access
     * @param name
     * @param descriptor
     * @param signature
     * @param exceptions
     * @return
     */
    @Override
    public MethodVisitor visitMethod(int access, String name, String descriptor, String signature, String[] exceptions) {
        MethodVisitor mv = super.visitMethod(access, name, descriptor, signature, exceptions);
        Method method = map.get(descriptor);
        if (method == null) {
            return mv;
        }
        return new MethodVisitor(Opcodes.ASM9, mv) {
            // 使用栈跟踪嵌套的 NEW 指令
            private final Deque<String> pendingNews = new ArrayDeque<>();

            @Override
            public void visitTypeInsn(int opcode, String type) {
                if (opcode == Opcodes.NEW) {
                    pendingNews.push(type); // 记录当前 NEW 的类型
                }
                super.visitTypeInsn(opcode, type);
            }

            @Override
            public void visitMethodInsn(int opcode, String owner, String name,
                                        String descriptor, boolean isInterface) {
                if (opcode == Opcodes.INVOKESPECIAL && !pendingNews.isEmpty()) {
                    String originalType = pendingNews.pop();

                    // 仅替换最外层的 NEW 指令（即返回值的类型）
                    if (pendingNews.isEmpty()) {
                        if (method != null) {
                            try {
                                Class orgType = Class.forName(originalType.replace('/', '.'), false, aop.getClassLoader());
                                Class enhanceType = aop.enhance(orgType);
                                if (enhanceType != orgType) {
                                    assert (orgType.isAssignableFrom(enhanceType));

                                    modified.add(method);
                                    String mappedType = AsmUtils.makeExternalName(enhanceType.getName());
                                    super.visitMethodInsn(opcode, mappedType, name, descriptor, isInterface);
                                    return;
                                }
                            } catch (ClassNotFoundException e) {
                                throw new ZoomException(e);
                            }
                        }
                    }
                }
                super.visitMethodInsn(opcode, owner, name, descriptor, isInterface);
            }

//            @Override
//            public void visitInsn(int opcode) {
//                // 处理 DUP 指令前的状态
//                if (opcode == Opcodes.DUP && !pendingNews.isEmpty()) {
//                    String type = pendingNews.peek();
//                    String mappedType = subclassMapping.getOrDefault(type, type);
//                    if (!type.equals(mappedType)) {
//                        // 替换 NEW 和 DUP 指令
//                        super.visitTypeInsn(Opcodes.NEW, mappedType);
//                        super.visitInsn(Opcodes.DUP);
//                        pendingNews.pop();
//                        pendingNews.push(mappedType);
//                        return;
//                    }
//                }
//                super.visitInsn(opcode);
//            }
        };
    }


}
